namespace il
{
	class CMAES
	{
		CMAES(int dimension, int populationSize=0)
		{			ndim = dimension;
			
			assert(ndim > 0);
					
			//Selection and Recombination
			sigma = .5;
			lambda = 10 + populationSize + int( 3 * log( 1.0 * ndim ) );

			mu = int(lambda/2.);
			
			double sum=0, sum2=0;
			
			for(int i = 0 ; i <  mu ; ++i)
				sum += log(lambda/2 + 0.5) - log(i+1.);

			omega.create( mu, mu );
			omega = 0;

			double testsum = 0;
			for(int i = 0 ; i <  mu ; ++i)
			{
				omega(i,i) = ( log(lambda/2 + 0.5) - log(i+1.) )/sum;
				testsum += omega(i,i);
				sum2 += omega(i,i)*omega(i,i);
			}

			assert(abs(testsum-1)<1e-10);
			
			mueff = 1./sum2;

			//Step-size control
			csigma = (mueff+2)/(ndim + mueff+5);
			double aux = sqrt((mueff-1)/(ndim+1))-1;
			aux = max<double>( 0., aux );
			dsigma = 1 + csigma + 2*aux;
			
			//Covariance matrix adaptation
			cc = ( 4.0 + mueff / ndim ) / ( ndim + 4.0 + 2.0 * mueff / ndim );
			c1 = 2.0 / ( ( ndim + 1.3 ) * ( ndim + 1.3 ) + mueff );
			aux = 2.0*( mueff - 2.0 + 1.0 / mueff ) / ( ( ndim + 2.0 ) * ( ndim + 2.0 ) + mueff );
			cmu = min<double>( 1.0 - c1 , aux );
			
			assert(cmu+c1 <= 1 );
			assert(cc <= 1);

			//Initialize dynamic strategy parameters and constants
			pc = Mat_<double>::zeros( ndim , 1 );
			ps = pc.clone();
			B =  Mat_<double>::eye( ndim , ndim );
			D = B.clone();
			
			C = B*D * (B*D).t();
			eigeneval = 0;
			chiN = sqrt( 1.0 * ndim ) * (1.0 - 1.0/(4.0*ndim) + 1.0/(21.0*ndim*ndim) );

			bestEnergy = numeric_limits<double>::max();
		}

		virtual ~CMAES(){}

		virtual double FitnessFunction( const Mat_<double> &x ) = 0;
			
		double energy(){return bestEnergy;}
		Mat_<double> solution(){return bestSolution;}

		void solve(const int &maxgen, const Mat_<double> &seed)
		{
			assert(maxgen!=3);
			cout << "CMA solver computing ...\n";
			long int to,tf;
			int generation = 0, counteval=0, stopeval = 1e5;
			//double stopenergy = 1.0e-10;
		
			//I'll need to define this struct with a operator overload for < so that 
			//I can use the std::list::sort() with no need to define complicated predicates
			struct energy_node
			{
				double val;
				int idx;
				bool operator<( energy_node b )
				{
					return this->val < b.val;
				}
			};
			
			list<struct energy_node> arenergy;
			Mat_<double> arx , arz, xsel , zsel, xmean ,zmean;

			arx.create( ndim , lambda );
			arz.create( ndim , lambda );
			xsel.create( ndim, mu );
			zsel.create( ndim, mu );
			xmean.create( ndim, 1 );
			zmean.create( ndim, 1 );
			
			if( seed.rows * seed.cols  )
				seed.copyTo( xmean );
			else
				randn( xmean, Scalar::all(0), Scalar::all(1.) );
			
			xmean.copyTo(bestSolution);

			bestEnergy = FitnessFunction( bestSolution );
	
			do //MAIN LOOP
			{
				to = getTickCount();
				//Generate and evaluate lambda offspring
				
				randn( arz, Scalar::all(0), Scalar::all(1) );
				arx = sigma * ( B * D * arz );
				
				for( int k = 0 ; k < lambda ; ++k )
				{
					arx.col(k) = arx.col(k) + xmean;
					struct energy_node node;
					node.val = FitnessFunction( arx.col(k) );
					node.idx = k;
					arenergy.push_back( node );
					counteval++;
				}
				
				//Sort by fitness and compute weighted mean
				
				arenergy.sort();
				list<struct energy_node>::iterator it = arenergy.begin();

				for(int counter = 0 ; counter < mu ; ++counter, ++it )
				{
					arx.col( it->idx ).copyTo( xsel.col(counter) );
					arz.col( it->idx ).copyTo( zsel.col(counter) );
				}
				
				xmean = xsel * omega.diag();
				zmean = zsel * omega.diag();

				//Cumulation: Update evolution paths
				ps = ( 1.0-csigma ) * ps + ( sqrt(cc * (2.0-cc) * mueff) ) * ( B * zmean );
				

	
				int hsigma = norm(ps) / ( sqrt( 1.0 - pow( 1.0-csigma, 2.0 * (generation+1) ))) / chiN <
					(1.4+2/(ndim+1)) ? 1:0;

				assert( hsigma==0 || hsigma==1 );
				
				pc = ( 1-cc ) * pc + hsigma * sqrt( cc * (2-cc) * mueff ) * ( B * D * zmean );
				
				//Adapt covariance matrix C
				C = ( 1.0-c1-cmu ) * C + c1 * ( pc * pc.t() + ( 1.0 - hsigma ) * cc * ( 2.0-cc ) * C ) + 
					cmu * (B*D*zsel) * omega * (B*D*zsel).t();

				//Adapt step-size sigma
				sigma *= exp( (csigma/dsigma)*( norm( ps )/chiN-1) );

				//Update B and D from C
				if( counteval - eigeneval > lambda/( (c1 + cmu)*ndim*10 ) )
				{
					//enforce symetry
					completeSymm( C );
				
					Mat_<double> diag2;
					
					eigen( C , diag2 , B );

					double condition_number = diag2(0,0)/diag2(ndim-1,0);

					if(condition_number > 1e10)	break;

					sqrt( diag2 , diag2 );
					D = Mat::diag( diag2 );

					B = B.t(); //put eigenvectors on columns
					for(int j = 0 ; j < B.cols ; ++j)
						normalize( B.row(j) , B.row(j) );

				}

				//Escape flat fitness, or better terminate??
				it = arenergy.begin();
				for( int i=0 ; i < int( 0.7 * lambda ) + 1 ; ++i )
					++it;

				if( arenergy.front().val == it->val )
				{
					sigma *= exp( 0.2 + csigma / dsigma );
					cout << "Warning: flat fitness, reformulating objective\n";
				}

				if( arenergy.front().val < bestEnergy )
				{
					bestEnergy = arenergy.front().val;
					xsel.col( 0 ).copyTo( bestSolution );
				}

				arenergy.clear();
				++generation;

			}while ( generation < maxgen );

		}
	
	protected:

		double bestEnergy, mueff, sigma,
				csigma, dsigma, cc, c1,
				cmu, chiN;

		Mat_<double> bestSolution;
		Mat_<double> pc, ps, B, D, C, omega;
		int ndim, lambda, mu, eigeneval;
	};
	class CMASolver
	{
	public:
		CMASolver(int dimension, int popSize=0)
		{
			ndim = dimension;
			assert(ndim > 0);
			_VERBOSE > 0 && cout << "Solution dimension: " << ndim << endl;
			
			//Selection and Recombination
			sigma = .5;
			lambda = 10 + popSize + int( 3 * log( 1.0 * ndim ) );

			_VERBOSE > 0 && cout << "Lambda: " << lambda << endl;

			mu = int(lambda/2.);
			
			double sum=0, sum2=0;
			
			for(int i = 0 ; i <  mu ; ++i)
				sum += log(lambda/2 + 0.5) - log(i+1.);

			omega.create( mu, mu );
			omega = 0;

			double testsum = 0;
			for(int i = 0 ; i <  mu ; ++i)
			{
				omega(i,i) = ( log(lambda/2 + 0.5) - log(i+1.) )/sum;
				testsum += omega(i,i);
				sum2 += omega(i,i)*omega(i,i);
			}

			assert(check_mat<double>(omega));
			assert(abs(testsum-1)<1e-10);
			
			mueff = 1./sum2;

			assert(check_number<double>(mueff));

			_VERBOSE > 1 && cout << "particles (lambda): "<< lambda << endl;
			_VERBOSE > 1 && cout << "mu: "<< mu << endl;
			_VERBOSE > 1 && cout << "mueff: " << mueff<<endl;

			//Step-size control
			csigma = (mueff+2)/(ndim + mueff+5);
			double aux = sqrt((mueff-1)/(ndim+1))-1;
			aux = max<double>( 0., aux );
			dsigma = 1 + csigma + 2*aux;
			
			_VERBOSE > 1 && cout << "csigma: " << csigma << " dsigma: " << dsigma;

			//Covariance matrix adaptation
			cc = ( 4.0 + mueff / ndim ) / ( ndim + 4.0 + 2.0 * mueff / ndim );
			c1 = 2.0 / ( ( ndim + 1.3 ) * ( ndim + 1.3 ) + mueff );
			aux = 2.0*( mueff - 2.0 + 1.0 / mueff ) / ( ( ndim + 2.0 ) * ( ndim + 2.0 ) + mueff );
			cmu = min<double>( 1.0 - c1 , aux );
			
			_VERBOSE > 1 && cout << "cc: " << cc << " c1: " << c1 << " cmu: " << cmu << endl;
			assert(cmu+c1 <= 1 );
			assert(cc <= 1);

			//Initialize dynamic strategy parameters and constants
			pc = Mat_<double>::zeros( ndim , 1 );
			ps = pc.clone();
			B =  Mat_<double>::eye( ndim , ndim );
			D = B.clone();
			
			C = B*D * (B*D).t();
			eigeneval = 0;
			chiN = sqrt( 1.0 * ndim ) * (1.0 - 1.0/(4.0*ndim) + 1.0/(21.0*ndim*ndim) );

			_VERBOSE > 2 && cout << "C: \n";
			_VERBOSE > 2 && print_mat<double>(C);
			_VERBOSE > 1 && cout << "chiN: " << chiN << endl;

			bestEnergy = numeric_limits<double>::max();
		}

		virtual ~CMASolver(){}

		virtual double energyFunction( const Mat_<double> &x ) = 0;
		virtual void generationCompleted() = 0;
		
		double energy(){return bestEnergy;}
		Mat_<double> solution(){return bestSolution;}

		void solve(const int &maxgen, const Mat_<double> &seed)
		{
			assert(maxgen!=3);
			_VERBOSE > 0 && cout << "CMA solver computing ...\n";
			long int to,tf;
			int generation = 0, counteval=0, stopeval = 1e5;
			//double stopenergy = 1.0e-10;
		
			//I'll need to define this struct with a operator overload for < so that 
			//I can use the std::list::sort() with no need to define complicated predicates
			struct energy_node
			{
				double val;
				int idx;
				bool operator<( energy_node b )
				{
					return this->val < b.val;
				}
			};
			
			list<struct energy_node> arenergy;
			Mat_<double> arx , arz, xsel , zsel, xmean ,zmean;

			arx.create( ndim , lambda );
			arz.create( ndim , lambda );
			xsel.create( ndim, mu );
			zsel.create( ndim, mu );
			xmean.create( ndim, 1 );
			zmean.create( ndim, 1 );
			
			if( seed.rows * seed.cols  )
				seed.copyTo( xmean );
			else
				randn( xmean, Scalar::all(0), Scalar::all(1.) );
			
			xmean.copyTo(bestSolution);
			
			to = getTickCount();
			bestEnergy = energyFunction( bestSolution );
			tf = getTickCount();

			_VERBOSE>0 && cout << "Time computing energy: " << (1.0*tf-to)/getTickFrequency() << endl;
			_VERBOSE>0 && cout << "Initial best:" << bestEnergy << endl;
			
			do //MAIN LOOP
			{
				to = getTickCount();
				//Generate and evaluate lambda offspring
				
				randn( arz, Scalar::all(0), Scalar::all(1) );
				arx = sigma * ( B * D * arz );
				
				for( int k = 0 ; k < lambda ; ++k )
				{
					arx.col(k) = arx.col(k) + xmean;
					struct energy_node node;
					node.val = energyFunction( arx.col(k) );
					node.idx = k;
					arenergy.push_back( node );
					counteval++;
				}
				
				//Sort by fitness and compute weighted mean
				
				arenergy.sort();
				list<struct energy_node>::iterator it = arenergy.begin();

				for(int counter = 0 ; counter < mu ; ++counter, ++it )
				{
					arx.col( it->idx ).copyTo( xsel.col(counter) );
					arz.col( it->idx ).copyTo( zsel.col(counter) );
				}
				
				xmean = xsel * omega.diag();
				zmean = zsel * omega.diag();

				_VERBOSE > 2 && cout << "xmean: \n";
				_VERBOSE > 2 && print_mat<double>( xmean.t() );
				
				assert( check_mat(xmean) );
				assert( check_mat(zmean) );

				//Cumulation: Update evolution paths
				ps = ( 1.0-csigma ) * ps + ( sqrt(cc * (2.0-cc) * mueff) ) * ( B * zmean );
				
				assert( check_mat<double>(ps) );

				_VERBOSE > 2 && cout << "ps:\n";
				_VERBOSE > 2 && print_mat<double>( ps.t() );
				_VERBOSE > 1 && cout << "norm(ps): " << norm( ps ) << endl;
	
				int hsigma = norm(ps) / ( sqrt( 1.0 - pow( 1.0-csigma, 2.0 * (generation+1) ))) / chiN <
					(1.4+2/(ndim+1)) ? 1:0;

				assert( hsigma==0 || hsigma==1 );
				_VERBOSE > 1 && cout << "hsigma: " << hsigma << endl;
				
				pc = ( 1-cc ) * pc + hsigma * sqrt( cc * (2-cc) * mueff ) * ( B * D * zmean );
				
				_VERBOSE > 2 && cout << "pc:\n";
				_VERBOSE > 2 && print_mat<double>( pc.t() );
				
				assert( check_mat<double>( pc ) );
				
				//Adapt covariance matrix C
				C = ( 1.0-c1-cmu ) * C + c1 * ( pc * pc.t() + ( 1.0 - hsigma ) * cc * ( 2.0-cc ) * C ) + 
					cmu * (B*D*zsel) * omega * (B*D*zsel).t();

				_VERBOSE > 2 && cout << "C: \n";
				_VERBOSE > 2 && print_mat<double>( C );

				//Adapt step-size sigma
				sigma *= exp( (csigma/dsigma)*( norm( ps )/chiN-1) );

				assert( check_number<double>( sigma ) );

				_VERBOSE > 2 && cout << "STEP-SIZE: " << sigma << " norm(ps): "<< norm( ps ) <<  endl; 
				
				//Update B and D from C
				if( counteval - eigeneval > lambda/( (c1 + cmu)*ndim*10 ) )
				{
					//enforce symetry
					completeSymm( C );
					_VERBOSE > 1 && cout << "C:\n";
					_VERBOSE > 1 && print_mat<double>( C );
					
					assert( check_mat<double>( C ) );
					
					Mat_<double> diag2;
					
					eigen( C , diag2 , B );
					_VERBOSE > 1 && cout << "diag2: \n";
					_VERBOSE > 1 && print_mat<double>( diag2.t() );
					double condition_number = diag2(0,0)/diag2(ndim-1,0);

					assert( check_mat<double>( diag2 ) );
					assert( check_mat<double>( B ) );
					assert( check_number<double>( condition_number ) );

					if(condition_number > 1e10)	break;

					_VERBOSE > 1 && cout << "COND NO: " << condition_number <<endl;

					sqrt( diag2 , diag2 );
					D = Mat::diag( diag2 );

					_VERBOSE > 2 && cout << "D: \n";
					_VERBOSE > 2 && print_mat<double>( D );

					B = B.t(); //put eigenvectors on columns
					for(int j = 0 ; j < B.cols ; ++j)
						normalize( B.row(j) , B.row(j) );

					_VERBOSE > 2 && cout << "B:\n";
					_VERBOSE > 2 && print_mat<double>( B );

					assert( check_mat<double>( D ) );
					assert( check_mat<double>( B ) );

				}

				////Break if fitness is good enough
				//if(arenergy.front().val <= stopenergy)
				//	break;

				//Escape flat fitness, or better terminate??
				it = arenergy.begin();
				for( int i=0 ; i < int( 0.7 * lambda ) + 1 ; ++i )
					++it;

				if( arenergy.front().val == it->val )
				{
					sigma *= exp( 0.2 + csigma / dsigma );
					cout << "Warning: flat fitness, reformulating objective\n";
				}

				if( arenergy.front().val < bestEnergy )
				{
					bestEnergy = arenergy.front().val;
					xsel.col( 0 ).copyTo( bestSolution );
				}

				tf = getTickCount();
				_VERBOSE > 0 && cout << "Generation: " << generation << "/" << maxgen << " -- " <<
					"Current Best Energy: " << bestEnergy  << " Execution Time: " << (tf - to)/getTickFrequency()
					<< endl;

				_VERBOSE > 1 && cout << "current sol: \n";
				_VERBOSE > 1 && print_mat<double>( bestSolution.t() );

				arenergy.clear();
				generationCompleted();
				++generation;

			}while ( generation < maxgen );

			assert( check_mat<double>(bestSolution)  );
			assert( check_number<double>(bestEnergy) );

			_VERBOSE > 0 && cout << "CMA completed\n";
		}
	
	protected:

		double bestEnergy, mueff, sigma,
				csigma, dsigma, cc, c1,
				cmu, chiN;

		Mat_<double> bestSolution;
		Mat_<double> pc, ps, B, D, C, omega;
		int ndim, lambda, mu, eigeneval;
		
	};

	class CMATest1: public CMASolver
	{
	public:
		CMATest1(int dim): CMASolver(dim,_POPULATION_SIZE){}
	
		double energyFunction(Mat_<double> x)
		{
			//Rastrigin
			double res = 0;
			for(int i = 0 ; i < ndim ; ++i)
				res += x(i,0) * x(i,0) - 10 * cos( 2 * CV_PI * x(i,0) );
			res += 10  *ndim;
			return res;
		}
	};
	
	class CMATest2: public CMASolver
	{
	public:
		CMATest2(int dim) : CMASolver(dim,_POPULATION_SIZE){}

		double energyFunction(Mat_<double> x)
		{
			//Rosenbrock
			double res = 0;
			for(int i = 0 ; i < ndim-1 ; ++i)
				res += pow(1-x(i,0),2) + 100.0*pow(x(i+1,0)-pow(x(i,0),2),2);
			return res;
		}
	};

	class CMATest3: public CMASolver
	{
	public:
		CMATest3(int dim) : CMASolver(dim,_POPULATION_SIZE){}

		double energyFunction(Mat_<double> x)
		{
			//Schwefel
			double res = 418.9829*ndim;
			for(int i = 0 ; i < ndim ; ++i)
				res -= x(i,0) * sin( sqrt(abs(x(i,0))) );
			return res;
		}
	};
}